import tensorflow as tf
import numpy as np
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt
import os
from logger import Logger
import matplotlib as mpl

class ModelTrainer:
    def __init__(self, model, save_dir='models'):
        self.model = model
        self.save_dir = save_dir
        self.history = None
        self.logger = Logger(name='model_trainer')
        os.makedirs(save_dir, exist_ok=True)
        
    def train(self, X, y, validation_split=0.2, batch_size=32, epochs=50, initial_epoch=0):
        """
        训练模型
        
        Args:
            X: 输入数据
            y: 标签数据
            validation_split: 验证集比例
            batch_size: 批次大小
            epochs: 训练轮数
            initial_epoch: 初始轮数（用于继续训练）
        """
        self.logger.info(f"开始训练模型 - 数据集大小: {len(X)}")
        self.logger.info(f"训练参数: batch_size={batch_size}, epochs={epochs}")
        
        # 计算训练集和验证集大小
        train_size = int(len(X) * (1 - validation_split))
        self.logger.info(f"训练集大小: {train_size}, 验证集大小: {len(X) - train_size}")
        
        # 优化回调函数
        callbacks = [
            # 早停策略
            tf.keras.callbacks.EarlyStopping(
                monitor='val_loss',
                patience=5,
                restore_best_weights=True
            ),
            # 模型检查点
            tf.keras.callbacks.ModelCheckpoint(
                'saved_models/model_latest.h5',
                save_best_only=True,
                monitor='val_loss'
            ),
            # 学习率调度器
            tf.keras.callbacks.ReduceLROnPlateau(
                monitor='val_loss',
                factor=0.2,
                patience=3,
                min_lr=1e-6,
                verbose=1
            ),
            # 添加TensorBoard支持
            tf.keras.callbacks.TensorBoard(
                log_dir='./logs',
                histogram_freq=1
            )
        ]
        
        # 使用优化后的回调函数列表
        self.history = self.model.fit(
            X, y,
            validation_split=validation_split,
            batch_size=batch_size,
            epochs=epochs,
            initial_epoch=initial_epoch,
            callbacks=callbacks
        )
        
        self.logger.info("模型训练完成")
        return self.history
    
    def evaluate(self, X_test, y_test):
        """评估模型"""
        test_loss, test_accuracy = self.model.evaluate(X_test, y_test)
        return {
            'test_loss': test_loss,
            'test_accuracy': test_accuracy
        }
    
    def predict(self, image):
        """预测单张图片"""
        # 确保图片维度正确 (1, height, width, channels)
        if len(image.shape) == 3:
            image = np.expand_dims(image, axis=0)
        
        predictions = self.model.predict(image)
        return predictions
    
    def plot_training_history(self):
        """绘制训练历史"""
        if self.history is None:
            print("没有训练历史可供绘制")
            return
            
        fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
        
        # 绘制准确率
        ax1.plot(self.history.history['accuracy'], label='训练准确率')
        ax1.plot(self.history.history['val_accuracy'], label='验证准确率')
        ax1.set_title('模型准确率')
        ax1.set_xlabel('轮次')
        ax1.set_ylabel('准确率')
        ax1.legend()
        
        # 绘制损失
        ax2.plot(self.history.history['loss'], label='训练损失')
        ax2.plot(self.history.history['val_loss'], label='验证损失')
        ax2.set_title('模型损失')
        ax2.set_xlabel('轮次')
        ax2.set_ylabel('损失')
        ax2.legend()
        
        plt.tight_layout()
        plt.show()
        
        # 在绘图之前添加这些配置
        plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # 对于 MacOS
        # 如果上面的字体不可用，可以尝试：
        # plt.rcParams['font.sans-serif'] = ['Heiti TC']  # 对于 MacOS
        plt.rcParams['axes.unicode_minus'] = False  # 正确显示负号
  